import json
from typing import Optional, List, Dict

from datasets import load_dataset


def text_to_messages(text: str) -> Optional[List[Dict[str, str]]]:
    messages = []
    for line in text.replace("\n\n", "\n").split("\n"):
        if line.strip().startswith("Human:"):
            if len(messages) == 0 or messages[-1]["role"] != "user":
                messages.append({'role': 'user', 'content': line.strip()[6:].strip()})
            else:
                messages[-1]['content'] += "\n" + line.strip()[6:].strip()
        elif line.strip().startswith("Assistant:"):
            if len(messages) == 0 or messages[-1]["role"] != "assistant":
                messages.append({'role': 'assistant', 'content': line.strip()[10:].strip()})
            else:
                messages[-1]['content'] += "\n" + line.strip()[10:].strip()
        elif len(messages) > 0:
            messages[-1]['content'] += "\n" + line.strip()

    return messages


def row_map(example):
    return {
        'chosen': text_to_messages(example['chosen']),
        'rejected': text_to_messages(example['rejected']),
    }

dataset = load_dataset("Anthropic/hh-rlhf")
dataset = dataset.map(row_map)

for split in ["train", "test"]:
    with open(f'hh_rlhf_en_{split}.jsonl', 'w', encoding='utf-8') as f:
        for row in dataset[split]:
            f.write(json.dumps(row) + '\n')

print("end")
